#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <iostream>
#include <cstdlib>
#include "common.h"
#include "test_kernel_launch.cuh" 
  __global__ void warm_up()
{

}
struct TestResult {
    std::string filename;
    float average_time;
};
template <typename T>
float test_case(std::vector<int> input_size, std::vector<int> output_size,std::string  filename) {
    CUDA_INIT();
    //预热
    for(int i = 0; i < 5; i++) {
        warm_up<<<1,1,0,stream>>>();
    }
    int in_batch = input_size[0];
    int in_height = input_size[1];
    int in_channels = input_size[2];
    int in_width = input_size[3];

    int out_batch = output_size[0];
    int out_height = output_size[1];
    int out_channels = output_size[2];
    int out_width = output_size[3];

    int in_elems = in_batch * in_height * in_channels * in_width;
    int out_elems = out_batch * out_height * out_channels * out_width;

    T *input = nullptr;
    T *output = nullptr;
    T *input_cpu = nullptr;
    T *output_cpu = nullptr;


    // 分配 CPU 内存
    input_cpu = (T *)malloc(sizeof(T) * in_elems);
    output_cpu = (T *)malloc(sizeof(T) * out_elems);

    // 分配 GPU 内存
    cudaMalloc((void **)&input, sizeof(T) * in_elems);
    cudaMalloc((void **)&output, sizeof(T) * out_elems);

    //input初始化
    for(int i = 0; i < in_batch * in_height * in_width * in_channels; i++)
    {
        input_cpu[i] = (i*7)%127;
    }

    // 将input数据从 CPU 拷贝到 GPU
    cudaMemcpy(input, input_cpu, sizeof(T) * in_elems, cudaMemcpyHostToDevice);

    //执行算子，记录时间
    
    float total_time = 0.0;
    int test_count = 1000;
    float average_time = 0.0;

    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    for (int i = 0; i < test_count; i++)
    {
        cudaEventRecord(start, 0);
        test_kernel_launch(input, output, in_batch, in_height, in_channels, in_width,out_batch ,out_height ,out_channels, out_width,in_elems,out_elems,stream);
        cudaEventRecord(stop, 0);
        cudaEventSynchronize(stop);

        float elapsed_time = 0.0f;
        cudaEventElapsedTime(&elapsed_time, start, stop);
        total_time += elapsed_time;
        
    }
    average_time = total_time / (test_count);
    printf("prop average time: %f ms\n",average_time);

    cudaDeviceSynchronize();
 
    // 将output数据从 GPU 拷贝回 CPU
    cudaMemcpy(output_cpu, output, sizeof(T) * out_elems, cudaMemcpyDeviceToHost);

    //保存输出，用于后续比对
    std::ofstream ori_file(filename, std::ios::binary);
    ori_file.write(reinterpret_cast<char*>(output_cpu), sizeof(T) * out_elems);
    ori_file.close();

    // 释放内存
    free(input_cpu);
    free(output_cpu);

    cudaFree(input);
    cudaFree(output);
    return average_time;
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cerr << "Usage: " << argv[0] << " <output_filename>" << std::endl;
        return 1;
    }

    std::string base_filename = argv[1];
    // std::string base_filename = "/home/sw/jxie/work/cuda2maca/more_all_test/ppl_cukernel_clip_ndarray/output_ori/";

    std::vector<TestResult> results;

    // Test case 1
    std::string filename1 =  base_filename + "/test1.dat";
    float average_time1 = test_case<half>({2,128,256,512},{2,128,256,512},filename1);
    results.push_back({filename1, average_time1});

    // Test case 2
    std::string filename2 =  base_filename + "/test2.dat";
    float average_time2 = test_case<half>({2,256,256,1024},{2,256,256,1024},filename2);
    results.push_back({filename2, average_time2});

    // Test case 3
    std::string filename3 =  base_filename + "/test3.dat";
    float average_time3 = test_case<half>({3,128,256,2048},{3,128,256,2048},filename3);
    results.push_back({filename3, average_time3});

    // 每个测试样例保留一个输出结果用于判断优化一致性
    std::ofstream json_file(base_filename + "/test_results.json");
    json_file << "[\n";
    for(size_t i = 0; i < results.size(); ++i) {
        json_file << "  {\n";
        json_file << "    \"filename\": \"" << results[i].filename << "\",\n";
        json_file << "    \"average_time\": " << results[i].average_time << "\n";
        json_file << "  }";
        if(i < results.size() - 1) json_file << ",";
        json_file << "\n";
    }
    json_file << "]\n";
    json_file.close();

    return 0;
    
}